import torch
from torch.nn import init
import torch.nn as nn
import random
import torchvision.transforms as transforms
#utils.py 是一个工具箱文件，包含了权重初始化、学习率调整、参数冻结、GAN损失、假图像池、图片保存等常用功能。


# Initialize the weight of the network
def init_weight(net, init_gain = 0.02):#net：要初始化的神经网络（nn.Module的子类）；init_gain：初始化时正态分布的标准差，默认是0.02。
    def init_func(m):#用于初始化单个层（module）的函数，会被应用到网络的每一层。
        classname = m.__class__.__name__#获取当前层的类名。
        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):#如果当前层有权重（weight），并且类名包含“Conv”或“Linear”。
            init.normal_(m.weight.data, 0.0, init_gain)#用均值为0，标准差为init_gain的正态分布初始化权重（m.weight.data）。
            if hasattr(m, 'bias') and m.bias is not None:#如果有 bias（偏置），且bias不为空：
                init.constant_(m.bias.data, 0.0)#则将其初始化为0。
        elif classname.find('BatchNorm2d') != -1:#如果当前层是 BatchNorm2d（二维批归一化层）：
            init.normal_(m.weight.data, 1.0, init_gain)#用均值为1，标准差为init_gain的正态分布初始化权重（m.weight.data）。
            init.constant_(m.bias.data, 0.0)#将偏置（bias）初始化为0。
    net.apply(init_func)#将init_func应用到网络的每一层。


# Define the lambda policy for the learning rate decay，定义学习率衰减策略。
class Lambda_rule():
    def __init__(self, start_epoch, initial_epoch, decay_epoch):
        self.start_epoch = start_epoch  #index of the first epoch，第一个epoch的索引。
        self.initial_epoch = initial_epoch #number of epochs with the initial learning rate，初始学习率持续的epoch数。
        self.decay_epoch = decay_epoch #number of epochs with learning rate decay，学习率衰减持续的epoch数。

    def step(self, epoch):
        return 1.0 - max(0, epoch + self.start_epoch - self.initial_epoch) / float(self.decay_epoch + 1)
    #这一行是核心逻辑，计算并返回一个缩放因子（通常用于乘以初始学习率，得到当前的实际学习率）。详细解释如下：
    #epoch + self.start_epoch：当前的实际 epoch（有些情况下训练可能不是从0开始，所以加上 start_epoch）。
    #epoch + self.start_epoch - self.initial_epoch：表示已经进入衰减阶段的 epoch 数（如果还没到 initial_epoch，这个值为负）。
    #max(0, ...)：如果还没到 initial_epoch，max 会让结果为0，保证前 initial_epoch 轮学习率不变。
    #float(self.decay_epoch + 1)：衰减阶段的总 epoch 数（+1 是为了防止分母为0，也让衰减更平滑）。
    #1.0 - ... / ...：随着 epoch 增加，分子变大，整个式子逐渐从1减小到0，实现线性衰减。

# Set if parameters of a network requires gradient
def set_requires_grad(nets, requires_grad = False):
    if not isinstance(nets, list):
        nets = [nets]
    # 如果 nets 不是列表，则将其转换为列表。
    for net in nets:
        if net is not None:#检查网络对象是否为 None，防止出错。
            for param in net.parameters():#遍历网络的参数。
                param.requires_grad = requires_grad#设置参数是否需要计算梯度。
                #requires_grad 参数决定了这个参数是否参与梯度计算。
                #如果 requires_grad 为 True，则参数会参与梯度计算，从而在反向传播时更新参数；
                #如果 requires_grad 为 False，则参数不会参与梯度计算，从而在反向传播时保持不变。
# 这个函数的作用是批量设置一个或多个神经网络的参数是否需要计算梯度，从而灵活控制哪些部分参与训练，哪些部分保持不变

# GAN loss for the network
class GANLoss(nn.Module):
    def __init__(self, target_real_label = 1.0, target_fake_label = 0.0):#初始化时可以指定“真实”标签和“假”标签的数值，默认真实为1.0，假为0.0。
        super(GANLoss, self).__init__()
        self.register_buffer('real_label', torch.tensor(target_real_label))
        self.register_buffer('fake_label', torch.tensor(target_fake_label))
        #register_buffer 会把变量注册为模型的“缓冲区”，它们会随着模型保存/加载而保存/加载，但不会被视为模型参数（不会被优化器更新）。
        self.loss = nn.MSELoss()#使用均方误差损失函数（MSELoss）计算损失。

    def get_target_tensor(self, prediction, target_is_real):
        # Return a tensor filled with ground-truth label, and has the same size as the prediction
        if target_is_real:
            target_tensor = self.real_label
        else:
            target_tensor = self.fake_label
        #如果 target_is_real 为真，返回与 prediction 形状相同、值全为 real_label 的张量；否则返回全为 fake_label 的张量。
        
        return target_tensor.expand_as(prediction)
        #expand_as(prediction) 让标签张量扩展成和判别器输出 prediction 一样的形状，便于逐元素计算损失。

    def __call__(self, prediction, target_is_real):
    #重载 __call__ 方法，使得 GANLoss 实例可以像函数一样被调用。
        target_tensor = self.get_target_tensor(prediction, target_is_real)#首先调用 get_target_tensor 得到目标标签张量。
        return self.loss(prediction, target_tensor)#然后用 MSELoss 计算 prediction（判别器输出）与目标标签之间的均方误差，作为损失返回。
#这个类的作用是：根据输入的判别器输出和目标（真实/假），自动生成对应的标签张量，并用 MSELoss 计算损失，适用于LSGAN。
#这样可以方便地在训练判别器和生成器时，分别传入不同的 target_is_real，自动得到正确的损失。

# Store and load previously generated fake images
# The implementation is in reference to https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix.git
#在训练生成对抗网络（GAN）时，维护一个“假图像池”，用于存储和管理历史生成的假图像。
#当生成器生成新的假图片时，会先检查图片池中是否已经有类似的图片，如果有，则直接使用池中的图片，而不是重新生成。
#这样可以减少计算量，提高训练效率。
class ImagePool():
    def __init__(self, pool_size = 50):
        self.pool_size = pool_size
        if self.pool_size > 0:
            # Create an empty pool，创建一个空池。
            self.num_imgs = 0#当前池中图片的数量。
            self.images = []#用于存储图片的列表。

    def query(self, images):
        # return an image from the image pool
        # If the pool size is 0, just return the input images，如果池大小为0，直接返回输入图片。
        if self.pool_size == 0:
            return images
        return_images = []#用于存储返回图片的列表。
        for image in images:
            image = torch.unsqueeze(image.data, 0)#在第0维度上增加一个维度，方便后续处理。
            if self.num_imgs < self.pool_size:
                # If the pool is not full, insert the current image，如果池未满（当前池中图片数量小于池大小），插入当前图片。
                self.num_imgs += 1#当前池中图片的数量加1。
                self.images.append(image)#将当前图片添加到池中。
                return_images.append(image)#将当前图片添加到返回图片列表中。
            else:
                p = random.uniform(0, 1)#生成一个0到1之间的随机数。
                if p > 0.5:#如果随机数大于0.5，则从池中随机选择一个图片，并插入当前图片。
                    # return a random image, and insert current image in the pool
                    random_id = random.randint(0, self.pool_size - 1)#随机选择一个池中图片的索引。
                    'Python 列表的索引是从 0 开始的。如果池子的大小是 50,合法的索引是 0 到 49,故上行-1。'
                    tmp = self.images[random_id].clone()#克隆池中图片。
                    self.images[random_id] = image#将当前图片插入到池中。
                    return_images.append(tmp)#将克隆的池中图片添加到返回图片列表中。
                else:
                    # return current image，如果随机数小于等于0.5，则直接返回当前图片。
                    return_images.append(image)#将当前图片添加到返回图片列表中。
        return_images = torch.cat(return_images, 0)#将返回图片列表中的图片拼接成一个张量。
        return return_images#返回拼接后的图片张量。

# Transform image tensor to png image
def save_image(tensor, name):
    unloader = transforms.ToPILImage()
    image = tensor.cpu().clone()  # we clone the tensor to not do changes on it
    image = image * 0.5 + 0.5
    image = image.squeeze(0)  # remove the fake batch dimension
    image = unloader(image)
    image.save(name, "PNG")